# 导入必要的库
import numpy as np  # 用于数值计算和处理数组
import pandas as pd  # 用于数据处理和分析，处理表格数据
import re  # 用于正则表达式操作，文本处理中常用
import string  # 包含字符串常量，如标点符号等



# 导入gensim库，用于自然语言处理任务
import gensim
from gensim.parsing.preprocessing import STOPWORDS  # 包含常用的停用词
from gensim.utils import simple_preprocess  # 简单的文本预处理函数

# 导入NLTK的词干提取器
from nltk.stem.porter import PorterStemmer

# 忽略警告信息，避免干扰输出
import warnings
warnings.filterwarnings("ignore")

# 导入数据集，只读取前1000行数据
medium = pd.read_csv('Medium_AggregatedData.csv', nrows = 1000)
# 查看数据集的前几行，了解数据的基本结构
medium.head()

# 数据筛选：去掉非英文文章，筛选点赞数多于等于25的文章
# 筛选出语言为英文的文章
medium = medium[medium['language'] == 'en']
# 筛选出总点赞数大于等于25的文章
medium = medium[medium['totalClapCount'] >= 25]

# 整理文章对应标签
def findTags(title):
    """
    根据文章标题查找对应的标签
    :param title: 文章标题
    :return: 该文章对应的标签列表
    """
    rows = medium[medium['title'] == title]  # 筛选出标题与输入标题相同的所有行
    tags = list(rows['tag_name'].values)  # 获取这些行中tag_name列的值，并转换为列表
    return tags

titles = medium['title'].unique()  # 获取所有文章的唯一标题

# 创建一个字典，用于存储文章标题和对应的标签
tag_dict = {'title': [], 'tags': []}

# 遍历所有文章标题，将标题和对应的标签添加到字典中
for title in titles:
    tag_dict['title'].append(title)
    tag_dict['tags'].append(findTags(title))

# 将字典转换为DataFrame
tag_df = pd.DataFrame(tag_dict)

# 添加标签到原始DataFrame中
# 去重，只保留每个标题的第一条记录
medium = medium.drop_duplicates(subset = 'title', keep = 'first')

def addTags(title):
    """
    根据文章标题添加对应的标签
    :param title: 文章标题
    :return: 该文章对应的标签列表，如果出现错误则返回NaN
    """
    try:
        tags = list(tag_df[tag_df['title'] == title]['tags'])[0]  # 获取对应标题的标签列表
    except:
        # 如果出现错误，假设没有标签，返回NaN
        tags = np.NaN
    return tags

# 将标签加入到原始DataFrame中，创建一个新的列'allTags'
medium['allTags'] = medium['title'].apply(addTags)

# 只保留需要的列
keep_cols = ['title', 'url', 'allTags', 'readingTime', 'author', 'text']
medium = medium[keep_cols]

# 去除标题为空的记录
null_title = medium[medium['title'].isna()].index  # 获取标题为空的记录的索引
medium.drop(index = null_title, inplace = True)  # 根据索引删除这些记录

# 重置索引
medium.reset_index(drop = True, inplace = True)

# 打印处理后数据集的形状
print(medium.shape)
# 查看处理后数据集的前几行
medium.head()

# 文本清洗（正则表达式）
def clean_text(text):
    """
    对输入的文本进行清洗操作
    :param text: 输入的文本
    :return: 清洗后的文本
    """
    # 使用原始字符串去掉http开头的链接
    text = re.sub(r'(?:(?:https?|ftp):\/\/)?[\w/\-?=%.]+\.[\w/\-?=%.]+', ' ', text)
    # 使用原始字符串去掉包含数字的单词
    text = re.sub(r'\w*\d\w*', ' ', text)
    # 使用原始字符串去掉标点符号并转换为小写
    text = re.sub(r'[%s]' % re.escape(string.punctuation), ' ', text.lower())
    # 去掉换行符
    text = text.replace('\n', ' ')
    # 去除多余的空格
    text = ' '.join(text.split())
    return text

# 对DataFrame中的'text'列应用清洗函数
medium['text'] = medium['text'].apply(clean_text)

# 一般都是用现成的停用词典，但是现成的往往难以满足自己的任务需求，还需要额外补充
# 可以自己添加，一个词一个词的加入，也可以基于统计方法来计算，比如词频最高的前100个词等
# 自己添加一部分停用词
stop_list = STOPWORDS.union(set(['data', 'ai', 'learning', 'time', 'machine', 'like', 'use', 'new', 'intelligence', 'need', "it's", 'way',
                                 'artificial', 'based', 'want', 'know', 'learn', "don't", 'things', 'lot', "let's", 'model', 'input',
                                 'output', 'train', 'training', 'trained', 'it', 'we', 'don', 'you', 'ce', 'hasn', 'sa', 'do', 'som',
                                 'can']))

# 将 stop_list 转换为 DataFrame
stop_df = pd.DataFrame(stop_list, columns=['stopword'])

# 使用 to_csv 方法保存为 CSV 文件
stop_df.to_csv('stopwords.csv', index=False)
# 去停用词
def remove_stopwords(text):
    """
    去除文本中的停用词
    :param text: 输入的文本
    :return: 去除停用词后的文本
    """
    clean_text = []
    for word in text.split(' '):
        if word not in stop_list and (len(word) > 2):  # 如果单词不在停用词列表中且长度大于2
            clean_text.append(word)
    return ' '.join(clean_text)

# 对DataFrame中的'text'列应用去除停用词函数
medium['text'] = medium['text'].apply(remove_stopwords)

# 词干提取
stemmer = PorterStemmer()

def stem_text(text):
    """
    对文本进行词干提取
    :param text: 输入的文本
    :return: 词干提取后的文本
    """
    word_list = []
    for word in text.split(' '):
        word_list.append(stemmer.stem(word))  # 对每个单词进行词干提取
    return ' '.join(word_list)

# 对DataFrame中的'text'列应用词干提取函数
medium['text'] = medium['text'].apply(stem_text)

# 预处理通常花的时间比较多，把结果保存下来
medium.to_csv('pre-processed.csv')
# 如果需要读取预处理后的数据，可以使用下面的代码
# medium = pd.read_csv('pre-processed.csv')